import time
from typing import Dict, List
import acl
import numpy as np
import os
import ctypes
from ctypes import c_void_p, c_int, c_size_t, c_ulong, c_int64,POINTER
ACL_MEM_MALLOC_HUGE_FIRST = 0
ACL_MEMCPY_HOST_TO_DEVICE = 1
ACL_MEMCPY_DEVICE_TO_HOST = 2
ACL_MEM_MALLOC_NORMAL_ONLY = 2
NPY_FLOAT32 = 11

libc = ctypes.CDLL("libc.so.6")
# mmap函数原型
mmap_func = libc.mmap
mmap_func.argtypes = [c_void_p, c_size_t, c_int, c_int, c_int, c_int64]
mmap_func.restype = c_void_p

# munmap函数原型
munmap_func = libc.munmap
munmap_func.argtypes = [c_void_p, c_size_t]
munmap_func.restype = c_int

def mmap_file(file_path):  
    # 打开文件并获取文件描述符  
    file_descriptor = os.open(file_path, os.O_RDONLY)  
    file_size = os.lseek(file_descriptor, 0, os.SEEK_END)  
    os.lseek(file_descriptor, 0, os.SEEK_SET)  
    # 调用mmap映射文件到内存  
    # PROT_READ和MAP_PRIVATE的值可能因系统而异，这里假设为1和2  
    protection_flags = 1  # PROT_READ  
    visibility_flags = 2  # MAP_PRIVATE  
    mapped_memory = mmap_func(None, file_size, protection_flags, visibility_flags, file_descriptor, 0)    
    if mapped_memory == -1:  
        raise Exception("Error mapping the file.")  

    # 关闭文件描述符，映射区域仍然有效  
    os.close(file_descriptor)  
    
    # 返回映射区域的地址  
    return mapped_memory,file_size

def check_ret(str,ret):
    if ret != 0:
        print(f"return code is {ret}, detail: {str}",flush=True) 

def initResource(device):
    ret = acl.init()
    check_ret("init", ret)
    ret = acl.rt.set_device(device)
    check_ret("set_device", ret)
    context,ret = acl.rt.create_context(device)
    check_ret("create_context", ret)
    return context

def destroyResource(device,context):
    ret = acl.rt.reset_device(device)
    ret = acl.finalize()
    ret = acl.rt.destroy_context(context)

dtype2NpType = {0:np.float32,1:np.float16,2:np.int8,3:np.int32,9:np.int64}

class ACLModel:
    def __init__(self,model_path,context=None,callback=None):
        self.context = context
        self.model_id = None
        self.model_desc = None
        self.callback_func = callback 
        self.tid = None
        self.stream = None
        self.callback_interval = 1
        self.exit_flag = False
        self.kvCache = None
        self.inputs_size_gear0 = [1*1025*8,1*1*8, 24*2*1*16*1024*128*2, 1*1*2048*2]
        self.inputs_size_gear1 = [1*1600*8,1*576*8, 24*2*1*16*1024*128*2, 1*576*2048*2]
        self.outputs_size_gear0 = [1*1*102400*4, 24*2*1*16*1*128*2]
        self.outputs_size_gear1 = [1*576*102400*4, 24*2*1*16*576*128*2]
        self.inputs_dataset_gear0, self.outputs_dataset_gear0 = None, None
        self.inputs_dataset_gear1, self.outputs_dataset_gear1 = None, None
        self.inputs_gear0:List[Dict[str,]] = []
        self.inputs_gear1:List[Dict[str,]] = []
        self.outputs_gear0:List[Dict[str,]] = []
        self.outputs_gear1:List[Dict[str,]] = []
        self.loadModel(model_path)
        self.past_key_values_index = 2
        self.allocateMem()
        if not callback:
            return
        self.stream, ret = acl.rt.create_stream()
        self.tid, ret = acl.util.start_thread(self._process_callback,
                                         [self.context, 50])
        check_ret("acl.util.start_thread", ret)
        ret = acl.rt.subscribe_report(self.tid, self.stream)
        check_ret("acl.rt.subscribe_report", ret)
    
    def unload(self):
        if self.callback_func:
            ret = acl.rt.synchronize_stream(self.stream)
            # 2.7 取消线程注册，Stream上的回调函数不再由指定线程处理。
            ret = acl.rt.unsubscribe_report(self.tid, self.stream)
            self.exit_flag = True
            ret = acl.util.stop_thread(self.tid)
            ret = acl.rt.destroy_stream(self.stream)
        self.freeMem()
        self.unloadModel()


    def loadModel(self, model_path):
        
        model_add, model_size = mmap_file(model_path)
        self.model_id, ret = acl.mdl.load_from_mem(model_add, model_size)
        
        #self.model_id, ret = acl.mdl.load_from_file(model_path)
        check_ret("load model",ret)
        munmap_func(model_add, model_size)
        self.model_desc = acl.mdl.create_desc()
        ret = acl.mdl.get_desc(self.model_desc, self.model_id)
        check_ret("get model desc",ret)
    
    def unloadModel(self):
        ret = acl.mdl.unload(self.model_id)
        if self.model_desc:
            ret  = acl.mdl.destroy_desc(self.model_desc)
            self.model_desc = None

    def allocateMem(self):
        self.inputs_dataset_gear0 = acl.mdl.create_dataset()
        self.inputs_dataset_gear1 = acl.mdl.create_dataset()
        self.inputs_gear0 = []
        self.inputs_gear1 = []
        for i in range(4):
            if i == self.past_key_values_index:
                buffer, ret = acl.rt.malloc(self.inputs_size_gear0[i], ACL_MEM_MALLOC_HUGE_FIRST)
                self.kvCache = acl.util.ptr_to_numpy(buffer,(24,2,1,16,1024,128),23)
                data = acl.create_data_buffer(buffer, self.inputs_size_gear0[i])
                _, ret = acl.mdl.add_dataset_buffer(self.inputs_dataset_gear0, data)
                check_ret("add_dataset_buffer",ret)
                _, ret = acl.mdl.add_dataset_buffer(self.inputs_dataset_gear1, data)                
                check_ret("add_dataset_buffer",ret)
                self.inputs_gear0.append({"buffer": buffer, "size": self.inputs_size_gear0[i]})
                self.inputs_gear1.append({"buffer": buffer, "size": self.inputs_size_gear0[i]})
            else:
                buffer0, ret = acl.rt.malloc(self.inputs_size_gear0[i], ACL_MEM_MALLOC_HUGE_FIRST)
                check_ret("alloc input memory",ret)
                data0 = acl.create_data_buffer(buffer0, self.inputs_size_gear0[i])
                _, ret = acl.mdl.add_dataset_buffer(self.inputs_dataset_gear0, data0)
                check_ret("add_dataset_buffer",ret)
                self.inputs_gear0.append({"buffer": buffer0, "size": self.inputs_size_gear0[i]})
                
                buffer1, ret = acl.rt.malloc(self.inputs_size_gear1[i], ACL_MEM_MALLOC_HUGE_FIRST)
                check_ret("alloc input memory",ret)                
                data1 = acl.create_data_buffer(buffer1, self.inputs_size_gear1[i])
                _, ret = acl.mdl.add_dataset_buffer(self.inputs_dataset_gear1, data1)
                check_ret("add_dataset_buffer",ret)
                self.inputs_gear1.append({"buffer": buffer1, "size": self.inputs_size_gear1[i]})
        index, ret = acl.mdl.get_input_index_by_name(self.model_desc, "ascend_mbatch_shape_data")
        dynamic_size = acl.mdl.get_input_size_by_index(self.model_desc,index)
        dynamic_buffer,ret = acl.rt.malloc(dynamic_size,ACL_MEM_MALLOC_HUGE_FIRST)
        check_ret("alloc input memory",ret)
        dynamic_data = acl.create_data_buffer(dynamic_buffer,dynamic_size)
        _, ret = acl.mdl.add_dataset_buffer(self.inputs_dataset_gear0, dynamic_data)
        check_ret("add_dataset_buffer",ret)
        _, ret = acl.mdl.add_dataset_buffer(self.inputs_dataset_gear1, dynamic_data)
        check_ret("add_dataset_buffer",ret)
         
        self.outputs_dataset_gear0 = acl.mdl.create_dataset()
        self.outputs_dataset_gear1 = acl.mdl.create_dataset()
        self.outputs_gear0 = []
        self.outputs_gear1 = []
        for i in range(2):
            data_type = acl.mdl.get_output_data_type(self.model_desc, i)
            
            buffer0, ret = acl.rt.malloc(self.outputs_size_gear0[i], ACL_MEM_MALLOC_HUGE_FIRST)
            check_ret("alloc output memory",ret)
            data0 = acl.create_data_buffer(buffer0, self.outputs_size_gear0[i])
            _, ret = acl.mdl.add_dataset_buffer(self.outputs_dataset_gear0, data0)
            check_ret("add_dataset_buffer",ret)
            buffer_host0, ret = acl.rt.malloc_host(self.outputs_size_gear0[i])
            check_ret("alloc output host memory",ret)
            self.outputs_gear0.append({"buffer": buffer0, "size": self.outputs_size_gear0[i],'buffer_host':buffer_host0,'dtype':dtype2NpType[data_type]})
            
            buffer1, ret = acl.rt.malloc(self.outputs_size_gear1[i], ACL_MEM_MALLOC_HUGE_FIRST)
            check_ret("alloc output memory",ret)
            data1 = acl.create_data_buffer(buffer1, self.outputs_size_gear1[i])
            _, ret = acl.mdl.add_dataset_buffer(self.outputs_dataset_gear1, data1)
            check_ret("add_dataset_buffer",ret)
            buffer_host1, ret = acl.rt.malloc_host(self.outputs_size_gear1[i])
            check_ret("alloc output host memory",ret)
            self.outputs_gear1.append({"buffer": buffer1, "size": self.outputs_size_gear1[i],'buffer_host':buffer_host1,'dtype':dtype2NpType[data_type]})           
            

    def freeMem(self):
        for item in self.input_data:
            ret = acl.rt.free(item["buffer"])
        ret = acl.mdl.destroy_dataset(self.input_dataset)
        for item in self.output_data:
            ret = acl.rt.free(item["buffer"])
            ret = acl.rt.free_host(item["buffer_host"])
        ret = acl.mdl.destroy_dataset(self.output_dataset)

    def inference(self,data,isPicture=False) -> List[np.ndarray]:
        #acl.rt.set_context(self.context)
        index, ret = acl.mdl.get_input_index_by_name(self.model_desc, "ascend_mbatch_shape_data")
        if isPicture:
            current_dims = {'name':'', 'dimCount':13,'dims':[1, 1600, 1, 576, 24, 2, 1, 16, 1024, 128, 1, 576, 2048]}
            ret = acl.mdl.set_input_dynamic_dims(self.model_id, self.inputs_dataset_gear1, index, current_dims)
            check_ret("set dims",ret)
        else : 
            current_dims = {'name':'', 'dimCount':13,'dims':[1, 1025, 1, 1, 24, 2, 1, 16, 1024, 128, 1, 1, 2048]}
            ret = acl.mdl.set_input_dynamic_dims(self.model_id, self.inputs_dataset_gear0, index, current_dims)             
            check_ret("set dims",ret)
        for i in range(len(data)):
            if i == self.past_key_values_index:
                pass
            else:
                bytes_data = data[i].tobytes()
                np_ptr = acl.util.bytes_to_ptr(bytes_data)
                if isPicture:
                    ret = acl.rt.memcpy(self.inputs_gear1[i]["buffer"], self.inputs_gear1[i]["size"], np_ptr,self.inputs_gear1[i]["size"], ACL_MEMCPY_HOST_TO_DEVICE)
                    check_ret("memcpy", ret)
                else:
                    ret = acl.rt.memcpy(self.inputs_gear0[i]["buffer"], self.inputs_gear0[i]["size"], np_ptr,self.inputs_gear0[i]["size"], ACL_MEMCPY_HOST_TO_DEVICE)
                    check_ret("memcpy", ret)
        if isPicture:
            ret = acl.mdl.execute(self.model_id, self.inputs_dataset_gear1,self.outputs_dataset_gear1)
            check_ret("execute ", ret)
        else :
            ret = acl.mdl.execute(self.model_id, self.inputs_dataset_gear0,self.outputs_dataset_gear0)
            check_ret("execute",ret)
                                
        inference_result = []
        if isPicture:
            for out in self.outputs_gear1:
                ret = acl.rt.memcpy(out['buffer_host'], out["size"],out["buffer"],out["size"],ACL_MEMCPY_DEVICE_TO_HOST)
                bytes_out = acl.util.ptr_to_bytes(out['buffer_host'], out["size"])
                out_data = np.frombuffer(bytes_out, dtype=out['dtype'])
                inference_result.append(out_data)
        else:
            for out in self.outputs_gear0:
                ret = acl.rt.memcpy(out['buffer_host'], out["size"],out["buffer"],out["size"],ACL_MEMCPY_DEVICE_TO_HOST)
                bytes_out = acl.util.ptr_to_bytes(out['buffer_host'], out["size"])
                out_data = np.frombuffer(bytes_out, dtype=out['dtype'])
                inference_result.append(out_data)
        return inference_result
    
    def inference_async(self,data,other_args) -> List[np.ndarray]:
        acl.rt.set_context(self.context)
        # print(f"wait lock {other_args[1]}",flush=True)
        # self.lock.acquire()
        # print(f"get lock {other_args[1]}",flush=True)
        for i in range(len(data)):
            bytes_data = data[i].tobytes()
            np_ptr = acl.util.bytes_to_ptr(bytes_data)
            ret = acl.rt.memcpy(self.inputs[i]["buffer"], self.inputs[i]["size"], np_ptr,self.inputs[i]["size"], ACL_MEMCPY_HOST_TO_DEVICE)
            check_ret("memcpy", ret)
        ret = acl.mdl.execute_async(self.model_id, self.input_dataset,self.output_dataset,self.stream)
        check_ret("exec_async", ret)
        print(f"submit exec task {other_args[1]}")
        ret = acl.rt.launch_callback(self.callPostProcess,other_args,1,self.stream)
        check_ret("launch callback", ret)

    def _process_callback(self, args_list):
        context, timeout = args_list
        acl.rt.set_context(context)
        while self.callback_interval:
            acl.rt.process_report(timeout)
            if self.exit_flag:
                print("[Callback] exit acl.rt.process_report")
                break

    def callPostProcess(self,other_args):
        print("start callback",flush=True)
        inference_result = []
        for out in self.outputs:
            ret = acl.rt.memcpy(out['buffer_host'], out["size"],out["buffer"],out["size"],ACL_MEMCPY_DEVICE_TO_HOST)
            bytes_out = acl.util.ptr_to_bytes(out['buffer_host'], out["size"])
            data = np.frombuffer(bytes_out, dtype=out['dtype'])
            inference_result.append(data)
        # self.lock.release()
        # print(f"free lock {other_args[1]}",flush=True)
        if not self.callback_func:
            return
        self.callback_func(inference_result,other_args)
        print(f"end callback, use time: {time.time()-time1}")
